#include <vector>
#include<iostream>
#include<stdio.h>
#include <algorithm>
#include <queue>
#include "non_max_suppression_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"
#include "helper.h"
using namespace std;

void print_arr(int32_t* array, int32_t arr_len){
    for(int32_t i=0; i<arr_len; ++i)
    {
        printf("%d ", array[i]);
        
    }
    printf("\n");
}


void bubbleSort(BoxInfoPtr data[], int n)//data[]是传过来的数组，n是数组中那些数的个数
{
    for (int i = 0; i < n - 1; i++)//这里是外循环，有n个数就比较n-1次
    {
        for (int j = 0; j<n - 1 - i; j++)//这是内循环，每轮比较n-1-（上一轮的除去的数的个数）
        {
            if (data[j].score_>data[j + 1].score_||(data[j].score_==data[j + 1].score_ && data[j].index_<data[j + 1].index_))//判断如果左边的数大于右边的数，就把大的数往右移
            {
                //这个是交换不用第三方变量的方式，当然也可以借用第三方变量来交换
                auto t = data[j];
                data[j] = data[j+1];
                data[j+1] = t;
            }
        }
    }
}


namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    TilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;

    // 1. 获取平台信息
    uint64_t ub_size;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto socVersion = ascendcPlatform.GetSocVersion();
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);

    
    const float* boxes_data = context->GetInputTensor(0)->GetData<float>();
    const float* scores_data = context->GetInputTensor(1)->GetData<float>();
    const int32_t max_output_boxes_per_class = (int32_t)*(context->GetInputTensor(2)->GetData<int32_t>());
    const float iou_threshold = *(context->GetInputTensor(3)->GetData<float>());
    const float score_threshold = *(context->GetInputTensor(4)->GetData<float>());
    
    const int32_t center_point_box = (int32_t)*(context->GetAttrs()->GetInt(0));

    auto boxes_data_shape = context->GetInputTensor(0)->GetOriginShape();
    auto boxes_tensor = context->GetInputTensor(0);
    auto scores_data_shape = context->GetInputTensor(1)->GetOriginShape();

    int32_t num_batches_ = boxes_data_shape.GetDim(0);
    int32_t num_classes_ = scores_data_shape.GetDim(1);
    int32_t num_boxes_ = boxes_data_shape.GetDim(1);

    // if (num_batches_*num_classes_>400) exit(-1);
    // if (num_boxes_>2000) exit(-1);
    uint32_t box_input_num = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    uint32_t score_input_num = context->GetInputShape(1)->GetStorageShape().GetShapeSize();
    uint32_t output_num = context->GetOutputShape(0)->GetStorageShape().GetShapeSize();

    tiling.set_max_output_boxes_per_class(max_output_boxes_per_class);
    tiling.set_iou_threshold(iou_threshold);
    tiling.set_score_threshold(score_threshold);
    tiling.set_center_point_box(center_point_box);
    tiling.set_num_batches_(num_batches_);
    tiling.set_num_classes_(num_classes_);
    tiling.set_num_boxes_(num_boxes_);
    tiling.set_box_input_num(box_input_num);
    tiling.set_score_input_num(score_input_num);
    tiling.set_output_num(output_num);





    context->SetBlockDim(1);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t* currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
static graphStatus InferDataType(gert::InferDataTypeContext* context)
{
    const auto inputDataType = context->GetInputDataType(0);
    context->SetOutputDataType(0, inputDataType);
    return ge::GRAPH_SUCCESS;
}
}

namespace ops {
class NonMaxSuppression : public OpDef {
public:
    explicit NonMaxSuppression(const char* name) : OpDef(name)
    {
        this->Input("boxes")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("scores")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("max_output_boxes_per_class")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT32})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
        
        this->Input("iou_threshold")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("score_threshold")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});
        
        this->Attr("center_point_box").AttrType(OPTIONAL).Int(0);

        this->Output("selected_indices")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT32})
                .Format({ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);
        this->AICore()
            .SetTiling(optiling::TilingFunc)
            .AddConfig("ascend310b");
    }
};
OP_ADD(NonMaxSuppression);
}
